from datasets import *


class Audio2ParamsNetwork(paddle.nn.Layer):
    def __init__(self, vgg16_pretrained=False):
        super().__init__()
        self.vgg = paddle.vision.models.vgg16(vgg16_pretrained)

    def forward(self, *inputs, **kwargs):

        return inputs


class AudioSimilarityNetwork(paddle.nn.Layer):
    def __init__(self, input_shape, pretrained=False):
        super(AudioSimilarityNetwork, self).__init__()
        self.vgg = paddle.vision.models.vgg16(pretrained)

        del self.vgg.avgpool
        del self.vgg.classifier

        flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
        self.fully_connect1 = paddle.nn.Linear(flat_shape, 512)
        self.fully_connect2 = paddle.nn.Linear(512, 1)

    def forward(self, x):
        x1, x2 = x
        x1 = self.vgg.features(x1)
        x2 = self.vgg.features(x2)
        b, _, _, _ = x1.size()
        x1 = x1.view([b, -1])
        x2 = x2.view([b, -1])
        x = paddle.abs(x1-x2)
        x = self.fully_connect1(x)
        x = self.fully_connect2(x)
        return x


class Piece2AudioNetwork(paddle.nn.Layer):
    def __init__(self, vgg16_pretrained=False):
        super().__init__()
        self.similarity_network = AudioSimilarityNetwork((224, 224, 3), vgg16_pretrained)
        self.fully_connect1 = paddle.nn.Linear(512, 1)

    def forward(self, *inputs, **kwargs):

        return inputs


def get_img_output_length(width, height):
    def get_output_length(input_length):
        # input_length += 6
        filter_sizes = [2, 2, 2, 2, 2]
        padding = [0, 0, 0, 0, 0]
        stride = 2
        for i in range(5):
            input_length = (input_length+2*padding[i]-filter_sizes[i]) // stride + 1
        return input_length
    return get_output_length(width)*get_output_length(height)


class EEG2ParamsNetwork(paddle.nn.Layer):
    def __init__(self):
        super().__init__()

    def forward(self, *inputs, **kwargs):
        return inputs


def audio2params_modelling(dataset_config):
    dataset_class = Audio2ParamsDataset
    train_dataset = dataset_class(mode='train', **dataset_config)
    test_dataset = dataset_class(mode='test', **dataset_config)
    network = Audio2ParamsNetwork()
    loss_fn = paddle.nn.MSELoss()
    return train_dataset, test_dataset, network, loss_fn


def piece2audio_modelling(dataset_config):
    dataset_class = Piece2AudioDataset
    train_dataset = dataset_class(mode='train', **dataset_config)
    test_dataset = dataset_class(mode='test', **dataset_config)
    network = Piece2AudioNetwork()
    loss_fn = paddle.nn.MSELoss()
    return train_dataset, test_dataset, network, loss_fn


def audio_similarity_modelling(dataset_config):
    dataset_class = Audio2AudioDataset
    train_dataset = dataset_class(mode='train', **dataset_config)
    test_dataset = dataset_class(mode='test', **dataset_config)
    network = AudioSimilarityNetwork()
    loss_fn = paddle.nn.MSELoss()
    return train_dataset, test_dataset, network, loss_fn


def eeg2midi_modelling(dataset_config):
    dataset_class = EEG2MIDIDataset
    train_dataset = dataset_class(mode='train', **dataset_config)
    test_dataset = dataset_class(mode='test', **dataset_config)
    network = EEG2ParamsNetwork()
    loss_fn = paddle.nn.CrossEntropyLoss()
    return train_dataset, test_dataset, network, loss_fn
